import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# 理想状态下 y = 4x + 2
x_data = [1.0, 2.0, 3.0, 4.0]
y_data = [6.0, 10.0, 14.0, 17.0]

def forward(x):
    '''

    :param x: 输入的参数
    :return: 返回一个y = w * x + b的函数计算式
    '''
    return w * x + b

def loss(x, y):
    '''
    该图中损失函数
    :param x:
    :param y:
    :return: 给定x所计算出来的y与理想的y之间的距离平方（保证距离大于0）
    '''
    y_pred = forward(x)
    return (y_pred - y) ** 2

w_list = []
b_list = []
mse_list = []

'''
    将w的数值在0-5.1之间，与b的数值在0-5.1之间的每个样本进行分别计算loss,
    之后将loss的总和除以样本量直接求出mes（Mean Square Error）
    通过Mes的值来进行对比程度
'''
for w in np.arange(0.0, 5.1, 0.1):
    for b in np.arange(0.0, 5.1, 0.1):
        print("w = ",w)
        print("b = ",b)
        l_sum = 0
        for x_val, y_val in zip(x_data, y_data):
            y_pred_val = forward(x_val)
            loss_val = loss(x_val, y_val)
            l_sum += loss_val
            print('\t', x_val, '\t', y_val, '\t', y_pred_val, '\t', loss_val)
        print('MES=', l_sum / 4)
        w_list.append(w)
        b_list.append(b)
        mse_list.append(l_sum / 4)

print('w_list=', w_list)
print('b_list=', b_list)
print('mse_list=', mse_list)

# 创建3D图形
fig = plt.figure()
ax = fig.add_subplot(projection='3d')

# 使用ax对象调用plot方法，而不是通过Axes3D类
ax.plot3D(w_list, b_list, mse_list, color='r')

# 添加轴标签使图形更清晰
ax.set_xlabel('w')
ax.set_ylabel('b')
ax.set_zlabel('MSE Loss')
ax.set_title('3D Loss Surface')

plt.show()



